#!/usr/bin/python3
# -*- coding:utf-8 -*-
# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
# This file is a part of the CANN Open Software.
# Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
# Please refer to the License for details. You may not use this file except in compliance with the License.
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
# See LICENSE in the root of the software repository for the full text of the License.
# ======================================================================================================================

import os
import torch
import numpy as np

def gen_golden_data_simple():
    dtype_f = np.float32
    dtype_i = np.int32

    # 输入
    low = np.iinfo(np.int32).min
    high = np.iinfo(np.int32).max

    current_scale = np.random.rand(1, 1).astype(dtype_f)
    growth_tracker = np.random.randint(low, high + 1, size=(1, 1), dtype=dtype_i)
    found_inf = np.random.rand(1, 1).astype(dtype_f)

    # 因子
    growthFactor = np.float32(np.random.rand())
    backoffFactor = np.float32(np.random.rand())
    growthInterval = np.random.randint(low, high + 1, dtype=np.int32)
    print(f"[INFO]growthFactor:{growthFactor}, backoffFactor:{backoffFactor}, growthInterval:{growthInterval}")

    # 输出
    updatedScale = 0.0
    updatedGrowthTracker = 0.0

    if (found_inf >= 1):
        current_scale *= backoffFactor
        growth_tracker = 0
    else:
        successful = growth_tracker + 1
        if (successful == growthInterval):
            current_scale *= growthFactor
            growth_tracker = 0
        else:
            growth_tracker = successful

    updatedScale = current_scale
    updatedGrowthTracker = growth_tracker

    # 将输入、cpu输出写入文件
    os.system("mkdir -p input")
    os.system("mkdir -p output")
    arr = np.array([[current_scale]], dtype=np.float32)
    updatedScaleT = torch.tensor(arr)
    arr = np.array([[growth_tracker]], dtype=np.int32)
    updatedGrowthTrackerT = torch.tensor(arr)

    updatedScaleT.numpy().astype(dtype_f).tofile("./output/golden_A.bin")
    updatedGrowthTrackerT.numpy().astype(dtype_i).tofile("./output/golden_B.bin")

    current_scale.astype(dtype_f).tofile("./input/input_x.bin")     # current_scale
    growth_tracker.astype(dtype_i).tofile("./input/input_y.bin")    # growth_tracker
    found_inf.astype(dtype_f).tofile("./input/input_z.bin")         # found_inf

    with open('./input/growthFactor.bin', 'w') as f:
        f.write(str(growthFactor) + '\n')
    with open('./input/backoffFactor.bin', 'w') as f:
        f.write(str(backoffFactor) + '\n')
    with open('./input/growthInterval.bin', 'w') as f:
        f.write(str(growthInterval) + '\n')

    print("[INFO]updatedScale:", updatedScale)
    print("[INFO]updatedGrowthTracker:", updatedGrowthTracker)

if __name__ == "__main__":
    gen_golden_data_simple()

